package com.example.springbootthreadexamples.thread;

import org.slf4j.MDC;
import org.springframework.core.task.TaskDecorator;

import java.util.Map;

/**
 * @author LanWP
 * @date 2024-08-02
 */
public class BizContextTaskDecorator implements TaskDecorator {
    @Override
    public Runnable decorate(Runnable runnable) {
        UserContextInfo userContext = UserContext.getUserContext();
        Map<String, String> copyOfContextMap = MDC.getCopyOfContextMap();
        return  () -> {
            UserContext.setUserContext(userContext);
            if (copyOfContextMap != null) {
                MDC.setContextMap(copyOfContextMap);
            }
            runnable.run();
        };
    }
}
